#####################################################################
# Functions
#####################################################################

# AR(1) with given error variance (based on a function in the replication package of Chernozhukov, Wuthrich, Zhu (2021, JASA))
generate.AR.series.var.input <- function(T01,rho,var.error){
  u     <- matrix(NA,T01,1)
  epsl  <- matrix(rnorm(T01),T01,1)*sqrt(var.error)
  startvalue <- rnorm(1)*sqrt(var.error/(1-rho^2))
  u[1,] <- rho*startvalue +epsl[1,]
  for (t in 2:T01){
    u[t,] <- rho*u[(t-1),]+epsl[t,]
  }
  return(u)
}

# AR(1)  (based on a function in the replication package of Chernozhukov, Wuthrich, Zhu (2021, JASA))
generate.AR.series <- function(T01,rho){
  u           <- matrix(NA,T01,1)
  epsl        <- matrix(rnorm(T01),T01,1)*sqrt(1-rho^2)
  startvalue  <- rnorm(1)
  u[1,]       <- rho*startvalue + epsl[1,] 
  for (t in 2:T01){
    u[t,] <- rho*u[(t-1),]+epsl[t,]
  }
  return(u)
}

# Oracle approach based on known w
oracle.cf <- function(y1,yJ,T0,T1,w.true,K){
  
  T01   <- T0 + T1
  r     <- min(floor(T0/K),T1)
  y1.pre  <- y1[1:T0]
  yJ.pre  <- yJ[1:T0,]
  y1.post <- y1[(T0+1):T01]
  yJ.post <- yJ[(T0+1):T01,]
  
  tau.mat <- matrix(NA,K,1)
  for (k in 1:K){
    Hk            <- (T0-(r*K))+seq((k-1)*r+1,k*r,1) 
    tau.mat[k,1]  <- mean(y1.post-yJ.post%*%w.true) - mean(y1.pre[Hk]-yJ.pre[Hk,]%*%w.true)
  }
  
  att <- mean(tau.mat)
  se  <- sqrt(1+((K*r)/T1))*sd(tau.mat)/sqrt(K)

  return(list(att=att,se=se))
}

# Implementation of Li (2020) based on Matlab code obtained from K. Li
li2020 <- function(Y1,Y0,T0,T1,m,nb,alpha.sig){
  
  T01 <- T0+T1
  y_pre   <- Y1[1:T0]
  y_post  <- Y1[(T0+1):(T01)]
  x_pre   <- Y0[1:T0,]
  x_post  <- Y0[(T0+1):(T01),]
  
  b_SC    <- as.matrix(scinference:::sc(y_pre,x_pre,lsei_type = 1)$w.hat)
  
  y_post_SC <- x_post%*%b_SC
  y_pre_SC <- x_pre%*%b_SC
  
  ATT_SC <- mean(y_post-y_post_SC)
  
  sigma_SC <- sqrt(mean((y_pre-y_pre_SC)^2))
  
  sigma2_v <- mean((y_post-y_post_SC-mean(y_post-y_post_SC))^2)
  
  e1_star <- sqrt(sigma2_v)*matrix(rnorm(T1*nb),T1,nb)
  
  A_star <- rep(NA,nb)
  for (g in 1:nb){
    xm        <- x_pre[(T0-m+1):T0,] # this is as in Li's code
    ym        <- xm%*%b_SC + sigma_SC*rnorm(m) 
    bm_SC     <- as.matrix(scinference:::sc(ym,xm,lsei_type = 1)$w.hat)
    A1_star   <- -mean(x_post%*%(bm_SC-b_SC))*sqrt(T1*m/T0)
    A2_star   <- sqrt(T1)*mean(e1_star[,g])
    A_star[g] <- A1_star + A2_star
    
  }
  
  ATT_order <- sort(A_star/sqrt(T1))
  
  cr_low    <- ATT_order[round(alpha.sig/2*nb)]
  cr_high   <- ATT_order[round((1-alpha.sig/2)*nb)]
  
  ci_low    <- ATT_SC - cr_high
  ci_high   <- ATT_SC - cr_low
  
  return(list(ATT_SC=ATT_SC,ci_low=ci_low,ci_high=ci_high))
  
}

# Permutation test as described in Abadie (2021, Section 3.5)
permutation <- function(Y1,Y0,T0,T1){
  
  J <- dim(Y0)[2]
  
  Y <- cbind(c(Y1),Y0)
  S.vec <- rep(NA,J+1)
  for (j in 1:(J+1)){
    Y1.temp <- Y[,j]
    Y0.temp <- Y[,-j]
    Y1.temp.pre <- Y1.temp[1:T0]
    Y0.temp.pre <- Y0.temp[(1:T0),]
    w.temp <- as.matrix(scinference:::sc(Y1.temp.pre,Y0.temp.pre,lsei_type=1)$w.hat) 
    u.hat <- Y1.temp - Y0.temp %*% w.temp
    S.vec[j] <- sqrt(mean((u.hat[(T0+1):(T0+T1)])^2))/sqrt(mean((u.hat[1:T0])^2))
  } 
  
  p.hat <- mean(S.vec>=S.vec[1])
  
  return(p.hat)
  
}

# One simulation draw
sim_one_sample <- function(DGP,T0,T1,J,K,Lambda,rho_u,var_u,var_factors,rho_vec,var_epsl_vec,w0_sc,tau_alternative,li,sdid,permtest){
  
  # Preliminaries
  sig_level <- 0.1
  
  # DGPs
  
  w_did     <- matrix(1,J,1)/J
  w_misspec <- c(-3,3,1,rep(0,J-3))
  
  if (DGP==1){ 
    w       <- w0_sc
    const   <- 0
    nonstat <- matrix(0,(T0+T1),J)
  }
  if (DGP==2){  
    w       <- w_did
    const   <- 0
    nonstat <- matrix(0,(T0+T1),J)
  }
  if (DGP==3){ 
    w       <- w_misspec
    const   <- 2
    nonstat <- matrix(0,(T0+T1),J)
  }
  if (DGP==4){ 
    w       <- w_misspec
    const   <- 2
    theta   <- c(1:(T0+T1))
    nonstat <- matrix(rep(theta,J),(T0+T1),J)
  }
  if (DGP==5){ 
    w       <- w_misspec
    const   <- 2
    theta   <- cumsum(rnorm((T0+T1)))
    nonstat <- matrix(rep(theta,J),(T0+T1),J)
  }
  if (DGP==6){ 
    w           <- w0_sc
    const       <- 0
    theta       <- c(1:(T0+T1))
    nonstat     <- matrix(rep(theta,J),(T0+T1),J)
    deviation   <- c(1:(T0+T1))
    nonstat[,1] <- nonstat[,1] + deviation 
  }
  if (DGP==7){ 
    w           <- w0_sc
    const       <- 0
    theta       <- cumsum(rnorm((T0+T1)))
    nonstat     <- matrix(rep(theta,J),(T0+T1),J)
    deviation   <- cumsum(rnorm((T0+T1)))
    nonstat[,1] <- nonstat[,1] + deviation 
  }
  if (DGP==8){ 
    w       <- w_misspec
    const   <- 0
    nonstat <- matrix(NA,(T0+T1),J)
    for (j in 1:J){
      nonstat[,j] <- j+j*c(1:(T0+T1))
    }
  }
  if (DGP==9){ 
    w       <- w0_sc
    const   <- 0
    nonstat <- matrix(NA,(T0+T1),J)
    for (j in 1:J){
      theta     <- rep(NA,(T0+T1))
      wn        <- rnorm((T0+T1))
      theta[1]  <- j + wn[1]
      for (t in 2:(T0+T1)){
        theta[t] <- j+theta[(t-1)] + wn[t]
      }
      nonstat[,j] <- theta
    }
  }
  
  T01 <- (T0+T1)
  
  u_sim <- generate.AR.series.var.input(T01,rho_u,var_u)
  
  Factors_sim <- matrix(NA,T01,4)
  for (i in 1:4){
    Factors_sim[,i] <- rnorm(T01)*sqrt(var_factors[i])
  }
  
  epsl_sim <- matrix(NA,T01,J)
  for (i in 1:J){
    epsl_sim[,i] <- generate.AR.series.var.input(T01,rho_vec[i],var_epsl_vec[i])
  }
  
  Y0 <- nonstat + Factors_sim %*% t(Lambda) + epsl_sim
  Y1 <- const + Y0 %*% w + u_sim 
  
  Y1[(T0+1):(T0+T1)] <- Y1[(T0+1):(T0+T1)]+tau_alternative
  
  # Results
  
  if (sdid == FALSE & li == FALSE){
    r_sc_sim    <- scinference(Y1,Y0,T1=T1,T0=T0,inference_method="ttest",estimation_method="sc",K=K,lsei_type = 1)
    tau_sc      <- r_sc_sim$att
    se_sc       <- r_sc_sim$se
    cov_sc      <- (abs(tau_sc/se_sc)<=qt(1-sig_level/2,df=K-1))
    leng_sc     <- 2*qt(1-sig_level/2,df=K-1)*se_sc
    
    r_did_sim     <- scinference(Y1,Y0,T1=T1,T0=T0,inference_method="ttest",estimation_method="did",K=K,lsei_type = 1)
    tau_did       <- r_did_sim$att
    se_did        <- r_did_sim$se
    cov_did       <- (abs(tau_did/se_did)<=qt(1-sig_level/2,df=K-1))
    leng_did      <- 2*qt(1-sig_level/2,df=K-1)*se_did
    
    w_true          <- w
    r_oracle_sim    <- oracle.cf(Y1,Y0,T0=T0,T1=T1,w_true,K)
    tau_oracle      <- r_oracle_sim$att
    se_oracle       <- r_oracle_sim$se
    cov_oracle      <- (abs(tau_oracle/se_oracle)<=qt(1-sig_level/2,df=K-1))
    leng_oracle     <- 2*qt(1-sig_level/2,df=K-1)*se_oracle
    
    leng_all <- c(leng_sc,leng_did,leng_oracle)
    cov_all  <- c(cov_sc,cov_did,cov_oracle)
    bias_all <- c(tau_sc,tau_did,tau_oracle)
  }
  if (li == TRUE){
    
    r_li2020  <- li2020(Y1,Y0,T0=T0,T1=T1,floor(T0*(2/3)),500,sig_level)
    leng_li   <- abs(r_li2020$ci_high-r_li2020$ci_low)
    cov_li    <- ((0 <= r_li2020$ci_high)*(0 >= r_li2020$ci_low)) 
    tau_li    <- r_li2020$ATT_SC
    
    leng_all <- leng_li
    cov_all  <- cov_li
    bias_all <- tau_li
    
  }
  
  if (sdid == TRUE){
    
    Ysdid         <- t(cbind(Y0,Y1))
    tau_sdid      <- synthdid_estimate(Ysdid, J, T0)
    se_sdid       <- sqrt(vcov(tau_sdid, method='placebo',replications = 100))
    cov_sdid      <- (abs(tau_sdid/se_sdid)<=qnorm(1-sig_level/2))
    leng_sdid     <- 2*qnorm(1-sig_level/2)*se_sdid
    
    leng_all <- leng_sdid
    cov_all  <- cov_sdid
    bias_all <- tau_sdid
    
  }
  
  if (permtest ==TRUE){
    cov_perm    <- c(permutation(Y1,Y0,T0=T0,T1=T1)>sig_level) 
    cov_all     <- c(cov_sc,cov_did,cov_oracle,cov_perm)
  }
  
  result  <- list(bias_all=bias_all,leng_all=leng_all,cov_all=cov_all)
  
  return(result)
  
}
